package com.oritsh.dicom.scu;

import com.oritsh.dicom.exception.DicomException;
import com.oritsh.dicom.exception.StoreException;
import org.dcm4che.data.Attributes;
import org.dcm4che.data.Tag;
import org.dcm4che.data.UID;
import org.dcm4che.io.DicomInputStream;
import org.dcm4che.net.*;
import org.dcm4che.net.pdu.AAssociateRQ;
import org.dcm4che.net.pdu.PresentationContext;

import java.io.IOException;
import java.util.*;

/**
 * Created with IntelliJ IDEA.
 * User: zarra
 * Date: 14-3-22
 * Time: 下午12:32
 * Copyright Shanghai Orient Rain Information Technology Co.,Ltd.
 */
public class StoreSCU extends DicomSCU {

    private Map<String, List<Attributes>> cache = new HashMap<>();

    private int priority;

    public StoreSCU() {
        super();
    }

    public StoreSCU(String aeCalling) {
        super(aeCalling);
    }


    public void addDicomInputStream(DicomInputStream dicomInputStream) throws DicomException {
        try {
            Attributes fmi = dicomInputStream.readFileMetaInformation();
            Attributes data = dicomInputStream.readDataset(-1, -1);
            String iuid = data.getString(Tag.SOPInstanceUID);
            cache.put(iuid, Arrays.asList(fmi, data));
        } catch (IOException e) {
            throw new StoreException(e.getMessage());
        }
    }

    @Override
    protected PresentationContext createPresentationContext() {
        PresentationContext presentationContext =
                new PresentationContext(1, UID.VerificationSOPClass, IVR_LE_FIRST);
        return presentationContext;
    }

    @Override
    protected AAssociateRQ createAAssociateRQ(ApplicationEntity aeCalling, ApplicationEntity aeCalled) {
        AAssociateRQ rq = super.createAAssociateRQ(aeCalling, aeCalled);

        for (Map.Entry<String, List<Attributes>> entry : cache.entrySet()) {
            Attributes fmi = entry.getValue().get(0);
            String ts = fmi.getString(Tag.TransferSyntaxUID);
            String cuid = fmi.getString(Tag.MediaStorageSOPClassUID);
            addPresentationContext(rq, cuid, ts);
        }
        //`确保能上传SecondaryCaptureImageStorage图像`
        String cuid = UID.SecondaryCaptureImageStorage;
        String ts = UID.ExplicitVRLittleEndian;
        addPresentationContext(rq, cuid, ts);
        return rq;
    }

    private void addPresentationContext(AAssociateRQ rq, String cuid, String tsuid) {
        if (!rq.containsPresentationContextFor(cuid)) {
            if (!tsuid.equals(UID.ImplicitVRLittleEndian))
                rq.addPresentationContext(
                        new PresentationContext(
                                rq.getNumberOfPresentationContexts() * 2 + 1,
                                cuid, UID.ImplicitVRLittleEndian));
            rq.addPresentationContext(
                    new PresentationContext(
                            rq.getNumberOfPresentationContexts() * 2 + 1,
                            cuid, tsuid));
        }
    }

    private void rawSend(String iuid, Attributes fmi, Attributes data, DimseRSPHandler handler) throws DicomException {
        String ts = fmi.getString(Tag.TransferSyntaxUID);
        String cuid = fmi.getString(Tag.MediaStorageSOPClassUID);
        try {
            as.cstore(cuid, iuid, priority, new DataWriterAdapter(data), ts,
                    handler);
        } catch (IOException | InterruptedException e) {
            throw new StoreException(e.getMessage());
        }

    }

    private DimseRSP rawSend(String iuid, Attributes fmi, Attributes data) throws DicomException {
        FutureDimseRSP rsp = new FutureDimseRSP(as.nextMessageID());
        rawSend(iuid, fmi, data, rsp);
        return rsp;
    }

    public void send(String iuid, DimseRSPHandler handler) throws DicomException {
        List<Attributes> value = cache.get(iuid);
        Attributes fmi = value.get(0);
        Attributes dataset = value.get(1);
        rawSend(iuid, fmi, dataset, handler);
    }

    public DimseRSP send(String iuid) throws DicomException {
        List<Attributes> value = cache.get(iuid);
        Attributes fmi = value.get(0);
        Attributes dataset = value.get(1);
        return rawSend(iuid, fmi, dataset);
    }

    public void sendAll(DimseRSPHandler handler) throws DicomException {
        Queue<String> queue = new LinkedList<>(cache.keySet());

        while (as.isReadyForDataTransfer()) {
            String iuid = queue.poll();
            if (iuid == null)
                break;
            send(iuid, handler);
        }
        try {
            as.waitForOutstandingRSP();
        } catch (InterruptedException e) {
            throw new StoreException(e.getMessage());
        }
    }

    public DimseRSP[] sendAll() throws DicomException {
        Queue<String> queue = new LinkedList<>();
        queue.addAll(cache.keySet());
        List<DimseRSP> rsps = new LinkedList<>();

        while (as.isReadyForDataTransfer()) {
            String iuid = queue.poll();
            if (iuid == null)
                break;
            DimseRSP rsp = send(iuid);
            rsps.add(rsp);
        }
        try {
            as.waitForOutstandingRSP();
        } catch (InterruptedException e) {
            throw new StoreException(e.getMessage());
        }

        return rsps.toArray(new DimseRSP[0]);
    }

    public List<Attributes> remove(String iuid){
        return cache.remove(iuid);
    }


}
